package org.offsync.apps

import java.io.IOException
import java.net.InetAddress
import java.net.InetSocketAddress
import java.net.Socket
import java.nio.ByteBuffer
import java.nio.channels.SelectionKey
import java.nio.channels.Selector
import java.nio.channels.ServerSocketChannel
import java.nio.channels.SocketChannel
import java.nio.channels.spi.SelectorProvider
import scala.util.control.Breaks

class NioServer(
    // The host:port combination to listen on
    private val hostAddress: InetAddress,
    private val port: Int,
    private val worker: EchoWorker) extends Runnable {
  // The channel on which we'll accept connections
  // Create a new selector
  private val selector = SelectorProvider.provider().openSelector()

  // Create a new non-blocking server socket channel
  val serverChannel = ServerSocketChannel.open()
  serverChannel.configureBlocking(false)

  // Bind the server socket to the specified address and port
  val isa = new InetSocketAddress(this.hostAddress, this.port)
  serverChannel.socket().bind(isa)

  // Register the server socket channel, indicating an interest in 
  // accepting new connections
  serverChannel.register(selector, SelectionKey.OP_ACCEPT)

  // The buffer into which we'll read data when it's available
  private val readBuffer: ByteBuffer = ByteBuffer.allocate(8192)

  // A list of PendingChange instances
  private val pendingChanges: java.util.List[ChangeRequest] = new java.util.LinkedList[ChangeRequest]()

  // Maps a SocketChannel to a list of ByteBuffer instances
  private val pendingData: java.util.Map[SocketChannel, java.util.List[ByteBuffer]] =
    new java.util.HashMap[SocketChannel, java.util.List[ByteBuffer]]()

  def send(socket: SocketChannel, data: Array[Byte]) {
    this.pendingChanges.synchronized {
      // Indicate we want the interest ops set changed
      this.pendingChanges.add(new ChangeRequest(socket, ChangeRequest.CHANGEOPS, SelectionKey.OP_WRITE))

      // And queue the data we want written
      this.pendingData.synchronized {
        var queue: java.util.List[ByteBuffer] = this.pendingData.get(socket)
        if (queue == null) {
          queue = new java.util.ArrayList[ByteBuffer]()
          this.pendingData.put(socket, queue)
        }
        
        queue.add(ByteBuffer.wrap(data))
      }
    }

    // Finally, wake up our selecting thread so it can make the required changes
    this.selector.wakeup()
  }

  def run() {
    while (true) {
      try {
        // Process any pending changes
        this.pendingChanges.synchronized {
          val changes = this.pendingChanges.iterator()
          
          while (changes.hasNext()) {
            val change = changes.next()
            
            change.`type` match {
            case ChangeRequest.CHANGEOPS =>
              val key = change.socket.keyFor(this.selector)
              key.interestOps(change.ops)
            }
          }
          
          this.pendingChanges.clear()
        }

        // Wait for an event one of the registered channels
        this.selector.select()

        // Iterate over the set of keys for which events are available
        val selectedKeys = this.selector.selectedKeys().iterator()
        while (selectedKeys.hasNext()) {
          val key: SelectionKey = selectedKeys.next()
          selectedKeys.remove()

          if (key.isValid()) {
            // Check what event is available and deal with it
            if (key.isAcceptable()) {
              this.accept(key)
            } else if (key.isReadable()) {
              this.read(key)
            } else if (key.isWritable()) {
              this.write(key)
            }
          }
        }
      } catch { case e: Exception =>
        e.printStackTrace()
      }
    }
  }

  private def accept(key: SelectionKey) = {
    // For an accept to be pending the channel must be a server socket channel.
    val serverSocketChannel = key.channel().asInstanceOf[ServerSocketChannel]

    // Accept the connection and make it non-blocking
    val socketChannel = serverSocketChannel.accept()
    val socket = socketChannel.socket()
    socketChannel.configureBlocking(false)

    // Register the new SocketChannel with our Selector, indicating
    // we'd like to be notified when there's data waiting to be read
    socketChannel.register(this.selector, SelectionKey.OP_READ)
  }

  private def read(key: SelectionKey): Unit = {
    val socketChannel = key.channel().asInstanceOf[SocketChannel]

    // Clear out our read buffer so it's ready for new data
    this.readBuffer.clear()

    // Attempt to read off the channel
    var numRead: Int = 0
    try {
      numRead = socketChannel.read(this.readBuffer)
    } catch { case e: IOException =>
      // The remote forcibly closed the connection, cancel
      // the selection key and close the channel.
      key.cancel()
      socketChannel.close()
      return
    }

    if (numRead == -1) {
      // Remote entity shut the socket down cleanly. Do the
      // same from our end and cancel the channel.
      key.channel().close()
      key.cancel()
      return
    }

    // Hand the data off to our worker thread
    this.worker.processData(this, socketChannel, this.readBuffer.array(), numRead)
  }

  private def write(key: SelectionKey) = {
    val socketChannel = key.channel().asInstanceOf[SocketChannel]

    this.pendingData.synchronized {
      val queue = this.pendingData.get(socketChannel)

      // Write until there's not more data ...
      val breaks = new Breaks
      breaks.breakable {
        while (!queue.isEmpty()) {
          val buf = queue.get(0)
          socketChannel.write(buf)
          if (buf.remaining() > 0) {
            // ... or the socket's buffer fills up
            breaks.break
          }
          queue.remove(0)
        }
      }

      if (queue.isEmpty()) {
        // We wrote away all data, so we're no longer interested
        // in writing on this socket. Switch back to waiting for
        // data.
        key.interestOps(SelectionKey.OP_READ)
      }
    }
  }
}

object NioServer {
  def main(args: Array[String]) = {
    try {
      val worker = new EchoWorker()
      new Thread(worker).start()
      new Thread(new NioServer(null, 9876, worker)).start()
    } catch { case e: IOException =>
      e.printStackTrace()
    }
  }
}
